from datetime import datetime, timedelta
from typing import Optional, Callable
from numpy import ndarray
from pandas import DataFrame
import re
import gm.api as gmdata
from vnpy.trader.setting import SETTINGS
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import BarData, TickData, HistoryRequest
from vnpy.trader.utility import round_to, ZoneInfo
from vnpy.trader.datafeed import BaseDatafeed

INTERVAL_VT2GM: dict[Interval, str] = {
    Interval.MINUTE: "60s",
    #Interval.HOUR: "60m",
    Interval.DAILY: "1d",
}

INTERVAL_ADJUSTMENT_MAP: dict[Interval, timedelta] = {
    Interval.MINUTE: timedelta(minutes=1),
    #Interval.HOUR: timedelta(hours=1),
    Interval.DAILY: timedelta() #no need to adjust for daily bar
}

FUTURES_EXCHANGES: set[Exchange] = {
    Exchange.CFFEX,
    Exchange.SHFE,
    Exchange.CZCE,
    Exchange.DCE,
    Exchange.INE,
    Exchange.GFEX,
}

EXCHANGE2GM: dict[Exchange, str] = {
    Exchange.CFFEX: 'CFFEX',
    Exchange.SHFE: 'SHFE',
    Exchange.CZCE: 'CZCE',
    Exchange.DCE: 'DCE',
    Exchange.INE: 'INE',
    Exchange.GFEX: 'GFEX',
}

CHINA_TZ = ZoneInfo("Asia/Shanghai")

def to_gm_symbol(symbol: str, exchange: Exchange) -> str :
    gm_symbol = EXCHANGE2GM.get(exchange, '') + '.' + symbol
    return gm_symbol

class GmdataDatafeed(BaseDatafeed) :
    def __init__(self):
        self.token: str = SETTINGS["datafeed.username"]
        self.inited: bool = False

    def init(self, output: Callable = print) -> bool :
        if self.inited:
            return True

        if not self.token :
            output("GMData数据服务初始化失败：token为空！")
            return False

        gmdata.set_token(self.token)

        self.inited = True
        return True

    def query_bar_history(self, req: HistoryRequest, output: Callable = print) -> Optional[list[BarData]] :
        if not self.inited :
            n: bool = self.init(output)
            if not n :
                return []
        ##############################################################
        symbol: str = req.symbol
        exchange: Exchange = req.exchange
        interval: Interval = req.interval
        start: datetime = req.start
        end: datetime = req.end
        ##############################################################
        if re.match('^[a-zA-Z]{1,}$', symbol) != None :
            symbol = symbol.upper()
        ##############################################################
        # 股票期权不添加交易所后缀
        if exchange in [Exchange.SSE, Exchange.SZSE] and symbol in self.symbols :
            gm_symbol: str = symbol
        else:
            gm_symbol: str = to_gm_symbol(symbol, exchange)

        gm_interval: str = INTERVAL_VT2GM.get(interval, None)
        if not gm_interval :
            output(f"GMData查询K线数据失败：不支持的时间周期{req.interval.value}")
            return []

        # 为了将GM时间戳（K线结束时点）转换为VeighNa时间戳（K线开始时点）
        adjustment: timedelta = INTERVAL_ADJUSTMENT_MAP[interval]

        # 只对衍生品合约才查询持仓量数据
        fields: list = ['open', 'high', 'low', 'close', 'volume', 'position', 'bob', 'eob']

        send_time = gmdata.get_next_n_trading_dates(exchange='SHFE', date=end.strftime('%Y-%m-%d'), n=1)
        send_time = send_time[0] + ' 18:00:00'
        df: DataFrame = gmdata.history(
            gm_symbol,
            frequency = gm_interval,
            start_time = start,
            end_time = send_time,
            fields = fields,
            adjust = gmdata.ADJUST_NONE,
            df = True,
        )

        data: list[BarData] = []

        if df is not None :
            #df.fillna(0, inplace=True)
            for row in df.itertuples():
                dt: datetime = row.bob.to_pydatetime()
                dt: datetime = dt.replace(tzinfo=CHINA_TZ)

                if dt >= end :
                    break

                bar: BarData = BarData(
                    symbol = req.symbol,
                    exchange=exchange,
                    interval=interval,
                    datetime=dt,
                    open_price=round_to(row.open, 0.000001),
                    high_price=round_to(row.high, 0.000001),
                    low_price=round_to(row.low, 0.000001),
                    close_price=round_to(row.close, 0.000001),
                    volume=row.volume,
                    open_interest=getattr(row, 'position', 0),
                    gateway_name='GM',
                )

                data.append(bar)

        return data

    def query_tick_history(self, req: HistoryRequest, output: Callable = print) -> Optional[list[TickData]] :
        if not self.inited :
            n: bool = self.init(output)
            if not n :
                return []

        symbol: str = req.symbol
        exchange: Exchange = req.exchange
        start: datetime = req.start
        end: datetime = req.end

        # 股票期权不添加交易所后缀
        if exchange in [Exchange.SSE, Exchange.SZSE] and symbol in self.symbols :
            gm_symbol: str = symbol
        else:
            gm_symbol: str = to_gm_symbol(symbol, exchange)

        # 只对衍生品合约才查询持仓量数据
        fields: list = [
            "open",
            "high",
            "low",
            "last",
            "prev_close",
            "volume",
            "total_turnover",
            "limit_up",
            "limit_down",
            "b1",
            "b2",
            "b3",
            "b4",
            "b5",
            "a1",
            "a2",
            "a3",
            "a4",
            "a5",
            "b1_v",
            "b2_v",
            "b3_v",
            "b4_v",
            "b5_v",
            "a1_v",
            "a2_v",
            "a3_v",
            "a4_v",
            "a5_v",
        ]
        if not symbol.isdigit() :
            fields.append("open_interest")

        df: DataFrame = gmdata.history(
            gm_symbol,
            frequency="tick",
            fields=fields,
            start_date=start,
            end_date=get_next_trading_date(end),        # 为了查询夜盘数据
            adjust_type="none"
        )

        data: list[TickData] = []

        if df is not None :
            # 填充NaN为0
            df.fillna(0, inplace=True)

            for row in df.itertuples():
                dt: datetime = row.Index[1].to_pydatetime()
                dt: datetime = dt.replace(tzinfo=CHINA_TZ)

                if dt >= end:
                    break

                tick: TickData = TickData(
                    symbol=symbol,
                    exchange=exchange,
                    datetime=dt,
                    open_price=row.open,
                    high_price=row.high,
                    low_price=row.low,
                    pre_close=row.prev_close,
                    last_price=row.last,
                    volume=row.volume,
                    turnover=row.total_turnover,
                    open_interest=getattr(row, "open_interest", 0),
                    limit_up=row.limit_up,
                    limit_down=row.limit_down,
                    bid_price_1=row.b1,
                    bid_price_2=row.b2,
                    bid_price_3=row.b3,
                    bid_price_4=row.b4,
                    bid_price_5=row.b5,
                    ask_price_1=row.a1,
                    ask_price_2=row.a2,
                    ask_price_3=row.a3,
                    ask_price_4=row.a4,
                    ask_price_5=row.a5,
                    bid_volume_1=row.b1_v,
                    bid_volume_2=row.b2_v,
                    bid_volume_3=row.b3_v,
                    bid_volume_4=row.b4_v,
                    bid_volume_5=row.b5_v,
                    ask_volume_1=row.a1_v,
                    ask_volume_2=row.a2_v,
                    ask_volume_3=row.a3_v,
                    ask_volume_4=row.a4_v,
                    ask_volume_5=row.a5_v,
                    gateway_name="GM"
                )

                data.append(tick)

        return data

